import paddle
from PIL import Image
from paddle.vision.datasets import DatasetFolder
from paddle.vision.transforms import transforms
from paddleslim.auto_compression import AutoCompression
import os

def get_transforms(image_size):
    normalize = transforms.Normalize(
        mean=[123.675, 116.28, 103.53], 
        std=[58.395, 57.120, 57.375]
    )
    return transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(image_size),
        transforms.Transpose(),
        normalize
    ])

class ImageNetDataset(DatasetFolder):
    def __init__(self, directory, image_size=224):
        super().__init__(directory)
        self.transform = get_transforms(image_size)
    
    def __getitem__(self, idx):
        img_path, _ = self.samples[idx]
        image = Image.open(img_path).convert('RGB')
        return self.transform(image)

    def __len__(self):
        return len(self.samples)

def get_dataloader(data_dir, image_size, batch_size):
    dataset = ImageNetDataset(data_dir, image_size)
    image = paddle.static.data(
        name='inputs', shape=[None, 3, image_size, image_size], dtype='float32'
    )
    loader = paddle.io.DataLoader(
        dataset, 
        feed_list=[image], 
        batch_size=batch_size, 
        shuffle=True, 
        drop_last=True,
        return_list=False
    )
    return loader

def main():
    paddle.enable_static()
    # 配置参数
    model_dir = "models"
    save_dir = "results"
    image_size = 224
    batch_size = 32
    train_data_dir = "datasets/ILSVRC2012/train/"
    model_filename = "inference.pdmodel"
    params_filename = "inference.pdiparams"

    train_loader = get_dataloader(train_data_dir, image_size, batch_size)
    
    # 自动压缩配置（根据您的操作系统修改 QuantPost/QuantAware 部分）
    config = {
        "QuantPost": {},
        "HyperParameterOptimization": {
            'ptq_algo': ['avg'],
            'max_quant_count': 3
        }
    }
    # config = {"QuantAware": {}, "Distillation": {}} # 如果为Windows请用此行（放开注释）

    # 启动模型压缩
    ac = AutoCompression(
        model_dir=model_dir,
        model_filename=model_filename,
        params_filename=params_filename,
        save_dir=save_dir,
        config=config,
        train_dataloader=train_loader,
        eval_dataloader=train_loader
    )
    ac.compress()

if __name__ == '__main__':
    main()